import dotenv
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_community.chat_models import ChatOpenAI
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate

dotenv.load_dotenv()

# 创建提示词
prompt = ChatPromptTemplate.from_messages(
    [
        ("system","你是一个很强大的聊天机器人，能够根据用户的上下文，来回复用户的问题，\n\n <context>{context}</context>"),
        ("human","{query}")
    ]
)

# 创建大语言模型
llm = ChatOpenAI(model="gpt-4o-mini",temperature=0)

# 创建链应用
chain = create_stuff_documents_chain(prompt=prompt,llm=llm)

# 文档列表
documents = [
    Document(page_content="小明喜欢绿色，但不喜欢黄色"),
    Document(page_content="小华喜欢粉色，有一点点喜欢红色"),
    Document(page_content="小章喜欢蓝色，但更喜欢青色"),
    Document(page_content="小hao喜欢蓝色，但更喜欢黄色")
]

# 调用链
content = chain.invoke({"query":"请帮我统计一下大家都喜欢什么颜色","context":documents})
print(content)